import os
import json
import glob
from datetime import datetime
import numpy as np

def convert_labelme_to_coco(labelme_dir, output_file):
    # 初始化COCO格式的数据结构
    coco_data = {
        "images": [],
        "type": "instances",
        "annotations": [],
        "categories": []
    }
    
    # 获取所有json文件
    json_files = glob.glob(os.path.join(labelme_dir, "*.json"))
    
    # 用于生成唯一ID
    image_id = 1
    annotation_id = 1
    
    # 处理每个json文件
    for json_file in json_files:
        with open(json_file, 'r', encoding='utf-8') as f:
            labelme_data = json.load(f)
        
        # print(labelme_data)
        # 获取图片信息
        image_info = {
            "id": image_id,
            "file_name": os.path.basename(labelme_data["imagePath"]),
            "width": labelme_data["imageWidth"],
            "height": labelme_data["imageHeight"]
        }
        coco_data["images"].append(image_info)
        
        # 处理标注信息
        for shape in labelme_data["shapes"]:
            # 获取边界框坐标
            points = np.array(shape["points"])
            x_min, y_min = points.min(axis=0)
            x_max, y_max = points.max(axis=0)
            width = x_max - x_min
            height = y_max - y_min
            
            # 创建分割信息（使用边界框的四个角点）
            segmentation = [
                [x_min, y_min, x_min, y_max, x_max, y_max, x_max, y_min]
            ]
            
            # 创建标注信息
            annotation = {
                "segmentation": segmentation,
                "area": width * height,
                "iscrowd": 0,
                "ignore": 0,
                "image_id": image_id,
                "bbox": [x_min, y_min, width, height],
                "category_id": int(shape["label"]),
                "id": annotation_id
            }
            coco_data["annotations"].append(annotation)
            annotation_id += 1
        
        image_id += 1
    
    # 添加类别信息
    categories = set()
    for shape in labelme_data["shapes"]:
        categories.add(int(shape["label"]))
    
    for category_id in sorted(categories):
        category = {
            "supercategory": "none",
            "id": category_id,
            "name": str(category_id)
        }
        coco_data["categories"].append(category)
    
    # 保存COCO格式文件
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(coco_data, f, indent=4, ensure_ascii=False)

if __name__ == "__main__":
    # 设置输入输出路径
    labelme_dir = "."  # 当前目录下的所有json文件
    output_file = "coco_annotations.json"
    
    # 执行转换
    convert_labelme_to_coco(labelme_dir, output_file)
    print(f"successfully! output_file:{output_file}") 